import base64
from Crypto.Cipher import PKCS1_v1_5 as PKCS1_cipher
from Crypto.PublicKey import RSA

default_encoding = "utf-8"


def bin_ord(flag):
    string = ""
    with open(flag) as f:
        txt = f.read()
        for i in range(len(txt)):
            string = string + bin(ord(txt[i])).replace('0b', '').zfill(8)
    return string


def get_key(key_or_path):
    if "BEGIN PUBLIC KEY" in key_or_path or "BEGIN PRIVATE KEY" in key_or_path:
        pem_data = key_or_path
    else:
        with open(key_or_path) as f:
            pem_data = f.read()
    return RSA.importKey(pem_data)


def rsa_encrypt(msg, pub_path, max_length=100):
    """
    RSA加密
    :param msg: 加密字符串
    :param pub_path: 公钥路径
    :param max_length: 1024bit的秘钥不能超过117， 2048bit的秘钥不能超过245
    :return:
    """
    key = get_key(pub_path)
    cipher = PKCS1_cipher.new(key)
    res_byte = bytes()
    for i in range(0, len(msg), max_length):
        res_byte += cipher.encrypt(msg[i:i + max_length].encode(default_encoding))
    return base64.b64encode(res_byte).decode(default_encoding)


def rsa_decrypt(msg, pri_path, max_length=256):
    """
    RSA解密
    :param msg: 加密字符串
    :param pri_path: 私钥路径
    :param max_length: 1024bit的秘钥用128，2048bit的秘钥用256位
    :return:
    """
    key = get_key(pri_path)
    cipher = PKCS1_cipher.new(key)

    res_bytes = bytes()
    encrypt_data = base64.b64decode(msg)
    for i in range(0, len(encrypt_data), max_length):
        res_bytes += cipher.decrypt(encrypt_data[i:i + max_length], 0)
    return res_bytes.decode(default_encoding)


if __name__ == '__main__':
    flag = r"/home/feifan/code/Python/ImageLSB_RSA/text/flag.txt"

    original_msg = bin_ord(flag)
    encrypted_data = rsa_encrypt(original_msg, "/home/feifan/code/Python/ImageLSB_RSA/key/pub_key.pem")
    print("encrypt_data:", encrypted_data)
    decrypted_data = rsa_decrypt(encrypted_data, "/home/feifan/code/Python/ImageLSB_RSA/key/pri_key.pem")
    ascii_text = ''.join([chr(int(decrypted_data[i:i+8], 2)) for i in range(0, len(decrypted_data), 8)])
    print("decrypt_data:", ascii_text)
